#include "StdAfx.h"
#include "FLVWriter.h"
#include "FFMpegInstance.h"

CFLVWriter::CFLVWriter(void)
{
	m_pFmtCtx = NULL;
	m_pOutFmt = NULL;
	m_pAudioSt = NULL;
	m_pVideoSt = NULL;

	m_bInitialized = FALSE;
	strcpy(m_fileName, "");
}

CFLVWriter::~CFLVWriter(void)
{
}

void CFLVWriter::Close()
{
	CFFMpegInstance::GetInstance()->f_av_write_trailer(m_pFmtCtx);
	CFFMpegInstance::GetInstance()->f_av_free(m_pFmtCtx);
	CFFMpegInstance::GetInstance()->f_av_free(m_pAudioSt);
	CFFMpegInstance::GetInstance()->f_av_free(m_pVideoSt);
}

BOOL CFLVWriter::Init(char fileName[], CodecID vCodecID, CodecID aCodecID, char* streamName)
{
	m_vCodecID = vCodecID;
	m_aCodecID = aCodecID;

	if (fileName == NULL || strlen(fileName) == 0)
		return FALSE;

	CFFMpegInstance::GetInstance()->f_lock();
	CFFMpegInstance::GetInstance()->f_avcodec_init();
	CFFMpegInstance::GetInstance()->f_avcodec_register_all();
	CFFMpegInstance::GetInstance()->f_av_register_all();

	m_pOutFmt = CFFMpegInstance::GetInstance()->f_guess_format("flv", fileName, NULL);

	if (m_pOutFmt == NULL)
	{
		CFFMpegInstance::GetInstance()->f_unlock();
		return FALSE;
	}

	m_pOutFmt->video_codec = vCodecID;
	m_pOutFmt->audio_codec = aCodecID;

	m_pFmtCtx = CFFMpegInstance::GetInstance()->f_av_alloc_format_context();

	if (m_pFmtCtx == NULL)
	{
		CFFMpegInstance::GetInstance()->f_unlock();
		return FALSE;
	}

	m_pFmtCtx->oformat = m_pOutFmt;
	strcpy(m_pFmtCtx->filename, fileName);

	m_pAudioSt = NULL;
	m_pVideoSt = NULL;

	if (m_pOutFmt->video_codec != CODEC_ID_NONE) 
	{
		m_pVideoSt = add_video_stream(m_pFmtCtx, m_pOutFmt->video_codec);
	}

	if (m_pOutFmt->audio_codec != CODEC_ID_NONE) 
	{
		m_pAudioSt = add_audio_stream(m_pFmtCtx, m_pOutFmt->audio_codec);
	}

	/* set the output parameters (must be done even if no
	parameters). */
	if (CFFMpegInstance::GetInstance()->f_av_set_parameters(m_pFmtCtx, NULL) < 0)
	{
		CFFMpegInstance::GetInstance()->f_unlock();
		return FALSE;
	}

	if(vCodecID != CODEC_ID_NONE)
	{
		CFFMpegInstance::GetInstance()->f_dump_format(m_pFmtCtx, m_pVideoSt->index, fileName, 1);
	}

	if (!(m_pOutFmt->flags & AVFMT_NOFILE))
	{
		if (CFFMpegInstance::GetInstance()->f_url_fopen(&m_pFmtCtx->pb, fileName, URL_WRONLY) < 0)
		{
			CFFMpegInstance::GetInstance()->f_unlock();
			return FALSE;
		}
	}

	CFFMpegInstance::GetInstance()->f_av_write_header(m_pFmtCtx);

	fn_insertStreamName(m_pFmtCtx->pb, streamName);

	CFFMpegInstance::GetInstance()->f_unlock();

	m_bInitialized = TRUE;
	return TRUE;
}

BOOL CFLVWriter::WriteVideo(unsigned char* lpData,  long biSizeImage, BOOL keyFrame)
{
	if(!m_bInitialized)
		return FALSE;
	
	if (lpData == NULL || biSizeImage == 0)
		return FALSE;

	int ret;
	AVCodecContext *c;
	static struct SwsContext *img_convert_ctx;

	c = m_pVideoSt->codec;

	AVPacket pkt;
	CFFMpegInstance::GetInstance()->f_av_init_packet(&pkt);

	//	if (c->coded_frame->pts != 0x8000000000000000)
	//        pkt.pts= CFFMpegInstance::GetInstance()->f_av_rescale_q(c->coded_frame->pts, c->time_base, m_pVideoSt->time_base);

	if(keyFrame)
		pkt.flags |= PKT_FLAG_KEY;

	pkt.stream_index = m_pVideoSt->index;
	pkt.data = (uint8_t *)lpData;
	pkt.size = biSizeImage;//sizeof(AVPicture);
	static int kuku = 1;
	//pkt.pts = kuku++;
	pkt.duration = (int)(1000.0f / (float)STREAM_FRAME_RATE);

	(m_pFmtCtx->streams[pkt.stream_index])->pts.num = 1;
	(m_pFmtCtx->streams[pkt.stream_index])->pts.den = STREAM_FRAME_RATE;

	ret = CFFMpegInstance::GetInstance()->f_av_write_frame(m_pFmtCtx, &pkt);

	if (ret != 0)
		return FALSE;
	
	return TRUE;
}

BOOL CFLVWriter::WriteAudio(unsigned char* lpData, long size)
{
	if(!m_bInitialized)
		return FALSE;
	
	if (lpData == NULL || size == 0)
		return FALSE;

	AVCodecContext *c;
	AVPacket pkt;
	CFFMpegInstance::GetInstance()->f_av_init_packet(&pkt);

	c = m_pAudioSt->codec;

	pkt.size= size;

	//    if (c->coded_frame->pts != 0x8000000000000000)
	//        pkt.pts= CFFMpegInstance::GetInstance()->f_av_rescale_q(c->coded_frame->pts, c->time_base, m_pAudioSt->time_base);

	pkt.flags |= PKT_FLAG_KEY;
	pkt.stream_index= m_pAudioSt->index;
	pkt.data= lpData;
	pkt.duration = (int)(1040.0f / (float)1);//STREAM_FRAME_RATE);
	//(m_pFmtCtx->streams[pkt.stream_index])->pts.num = 1;
	//(m_pFmtCtx->streams[pkt.stream_index])->pts.den = 5;

	/* write the compressed frame in the media file */
	if (CFFMpegInstance::GetInstance()->f_av_write_frame(m_pFmtCtx, &pkt) != 0) 
	{
		return FALSE;
	}

	return TRUE;
}

AVStream * CFLVWriter::add_video_stream(AVFormatContext *oc, CodecID codec_id)
{
	AVCodecContext *c;
	AVStream *st;

	st = CFFMpegInstance::GetInstance()->f_av_new_stream(oc, 0);

	if (!st) 
	{
		return NULL;
	}

	st->sample_aspect_ratio.num = 0;
	st->sample_aspect_ratio.den = 1;
	//st->sample_aspect_ratio = (AVRational){0,1};
	st->time_base.num = 1;
	st->time_base.den = STREAM_FRAME_RATE;

	AVCodec* codec = CFFMpegInstance::GetInstance()->f_avcodec_find_encoder((CodecID)codec_id);

	c = st->codec;
	c->codec_id = codec_id;
	c->codec_type = CODEC_TYPE_VIDEO;
	c->codec = codec;

	//if (CFFMpegInstance::GetInstance()->f_avcodec_open(c, codec) < 0)
	//	return NULL;

	/* put sample parameters */
	
	/* resolution must be a multiple of two */
	c->width = 320;
	c->height = 240;
	c->bit_rate = (c->width * c->height * 
		(((((c->height + c->width)/ 100)- 5)>> 1)+ 10) * 80 ) / 100;
	if (c->bit_rate < 300000) c->bit_rate = 300000;
	c->sample_aspect_ratio.num = 0;
	c->sample_aspect_ratio.den = 1;
	//c->sample_aspect_ratio = (AVRational){0,1};
	c->gop_size = 12;

	/* time base: this is the fundamental unit of time (in seconds) in terms
	of which frame timestamps are represented. for fixed-fps content,
	timebase should be 1/framerate and timestamp increments should be
	identically 1. */
	c->time_base.den = STREAM_FRAME_RATE;
	c->time_base.num = 1;
	c->max_b_frames =  0;
	c->pix_fmt = PIX_FMT_YUV420P;

	if (c->codec_id == CODEC_ID_MPEG2VIDEO) 
	{
		/* just for testing, we also add B frames */
		c->max_b_frames = 2;
	}
	
	if (c->codec_id == CODEC_ID_MPEG1VIDEO)
	{
		/* Needed to avoid using macroblocks in which some coeffs overflow.
		This does not happen with normal video, it just happens here as
		the motion of the chroma plane does not match the luma plane. */
		c->mb_decision=2;
	}

	// some formats want stream headers to be separate
	if(!strcmp(oc->oformat->name, "mp4") || !strcmp(oc->oformat->name, "mov") || !strcmp(oc->oformat->name, "3gp"))
		c->flags |= CODEC_FLAG_GLOBAL_HEADER;

	return st;
}

AVStream *CFLVWriter::add_audio_stream(AVFormatContext *oc, CodecID codec_id)
{
	AVCodecContext *c;
	AVStream *st;

	st = CFFMpegInstance::GetInstance()->f_av_new_stream(oc, 1);

	if (!st) 
	{
		return NULL;
	}

	AVCodec* codec = CFFMpegInstance::GetInstance()->f_avcodec_find_encoder((CodecID)codec_id);

	c = st->codec;
	c->codec_id = codec_id;
	c->codec_type = CODEC_TYPE_AUDIO;
	c->codec = codec;

	/* put sample parameters */
	c->bit_rate = 0;
	c->sample_rate = 11025;
	c->channels = 1;
	c->time_base.num = 1;
	c->time_base.den = 100;

	//c->bits_per_raw_sample = 8;
	//c->bits_per_coded_sample = 8;
	c->sample_fmt = SAMPLE_FMT_S16;
	//c->block_align = 1;
	//c->time_base.num = 1;

	return st;
}

void CFLVWriter::fn_insertStreamName(ByteIOContext *pb, char* streamName)
{
	int nameLength;
	nameLength = strlen(streamName);

	char tmpBuffer[3000];
	memset(tmpBuffer, 0, 3000);
	memcpy(tmpBuffer, pb->buffer, 15);

	int length = 0;
	unsigned char *pLen = (unsigned char*)&length;
	memcpy(pLen, pb->buffer + 16, 1);
	memcpy(pLen + 1, pb->buffer + 15, 1);
	length += 15;
	length += nameLength;
	memcpy(tmpBuffer + 15, pLen + 1, 1);
	memcpy(tmpBuffer + 16, pLen, 1);

	memcpy(tmpBuffer + 17, pb->buffer + 17, 229);

	char tmpByte4[1] = {0x0c};
	memcpy(tmpBuffer + 41, tmpByte4, 1);

	char tmpByte0[2] = {0x00, 0x0a};
	memcpy(tmpBuffer + 246, tmpByte0, 2);
	memcpy(tmpBuffer + 248, "streamName", 10);
	char tmpByte1[1] = {0x02};
	memcpy(tmpBuffer + 258, tmpByte1, 1);

	unsigned char *pNameLen = (unsigned char*)&nameLength;
	memcpy(tmpBuffer + 259, pNameLen + 1, 1);
	memcpy(tmpBuffer + 260, pNameLen, 1);

	memcpy(tmpBuffer + 261, streamName, nameLength);

	memcpy(tmpBuffer + 261 + nameLength, pb->buffer + 246, 300);

	memcpy(pb->buffer, tmpBuffer, 400);
	pb->buf_ptr = pb->buf_ptr + 15 + nameLength;
}
